#include <stdint.h>
#include <stddef.h>
#include <string.h>
#include <stdlib.h>
#include <stdio.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <unistd.h>
#include <stdbool.h>
#include "vmsa.h"

/**
 * @brief  read the exact number of bytes from file
 *
 * @param  [in]  filename     The file name
 * @param  [in]  buffer       The buffer to store the data
 * @param  [in]  len          The length of data to be read
 *
 * @returns 0 if the data read out equals to the required length
 *          -1 otherwise
 */
int read_file(char *filename, unsigned char *buffer, int len)
{
    FILE    *fp        = NULL;
    size_t  bytes_read = 0;

    if (!filename || !buffer) {
       printf("invalid file name or invalid buffer \n");
       return -1;
    }

   fp = fopen(filename, "rb");

    if (fp == NULL) {
        printf("Failed to open file %s\n", filename);
        return -1;
    }

    bytes_read = fread(buffer, 1, len, fp);

    fclose(fp);

    if (bytes_read < (size_t)len) {
        return -1;
    }

    return 0;
}

/**
 * @brief  write the exact number of bytes to file
 *
 * @param  [in]  filename     The file name
 * @param  [in]  buffer       The data buffer to write to file
 * @param  [in]  len          The length of data
 *
 * @returns 0 if the data read out equals to the required length
 *          -1 otherwise
 */
int write_file(char *filename, unsigned char *buffer, int len)
{
    FILE    *fp         = NULL;
    size_t  bytes_write = 0;

    if (!filename || !buffer) {
       printf("invalid file name or invalid buffer \n");
       return -1;
    }

   fp = fopen(filename, "wb");

    if (fp == NULL) {
        printf("Failed to open file %s\n", filename);
        return -1;
    }

    bytes_write = fwrite(buffer, 1, len, fp);

    fclose(fp);

    if (bytes_write < (size_t)len) {
        return -1;
    }

    return 0;
}

/**
 * @brief  get the file length
 *
 * @param  [in]  filename     The file name
 *
 * @returns  0 if file doesnot exist otherwise real length
 */
uint32_t  get_file_length(const char *filename)
{
    struct stat file_stat;

    if (filename == NULL || stat(filename, &file_stat) == -1) {
        return 0;
    }
    return file_stat.st_size;
}

/**
 * @brief  check if the file exists
 *
 * @param  [in]  filename     The file name
 *
 * @returns  true if file exists otherwise false
 */
uint32_t  file_exist(const char *filename)
{
    if (access(filename, F_OK) != -1) {
        return true;
    } else {
        return false;
    }
}

void cpuid(int function, int *eax, int *ebx, int *ecx, int *edx) {
    __asm__ volatile (
        "cpuid"
        : "=a" (*eax), "=b" (*ebx), "=c" (*ecx), "=d" (*edx)
        : "a" (function)
    );
}

struct x86_vcpu_model {
    const char *name;
    int family;
    int model;
    int stepping;
};

static const struct x86_vcpu_model builtin_x86_model[] = {
    {
        .name = "Dhyana",
        .family = 24,
        .model = 0,
        .stepping = 1,
    },
    {
        .name = "Dhyana-v1",
        .family = 24,
        .model = 0,
        .stepping = 1,
    },
    {
        .name = "Dhyana-v2",
        .family = 24,
        .model = 0,
        .stepping = 1,
    },
    {
        .name = "Dhyana-v3",
        .family = 24,
        .model = 0,
        .stepping = 1,
    },
    {
        .name = "Dharma",
        .family = 24,
        .model = 4,
        .stepping = 0,
    },
    {
        .name = "Dharma-v1",
        .family = 24,
        .model = 4,
        .stepping = 0,
    },
};

#define  EXT_FAMILY_MASK     0xFF00000
#define  EXT_MODEL_MASK      0xF0000
#define  BASE_FAMILY_MASK    0xF00
#define  BASE_MODEL_MASK     0xF0
#define  STEPPING_MASK       0xF

struct x86_vcpu_model  host_model;
static const struct x86_vcpu_model *get_vcpu_model(const char* name)
{
    int i;
    const struct x86_vcpu_model *model;

    if (name == NULL)
        return NULL;

    for (i = 0; i < sizeof(builtin_x86_model)/sizeof(builtin_x86_model[0]); i++) {
         model = &builtin_x86_model[i];

         if (strcmp(model->name, name) == 0) {
             return model;
         }
    }

    /* if vcpu_model is 'host' */
    if (strcmp(name, "host") == 0) {
        int eax, ebx, ecx, edx;
        uint8_t ext_family, base_family, ext_model, base_model;

        /* APM volumn 3
         * E.3.2 Function 1h—Processor and Processor Feature Identifiers
         */
        cpuid(1, &eax, &ebx, &ecx, &edx);
        ext_family = (eax & EXT_FAMILY_MASK) >> 20;
        ext_model   = (eax & EXT_MODEL_MASK) >> 16;
        base_family = (eax & BASE_FAMILY_MASK) >> 8;
        base_model  = (eax & BASE_MODEL_MASK) >> 4;

        if (base_family < 0xF) {
            host_model.family = base_family;
        } else {
            host_model.family =  base_family + ext_family;
        }

        if (base_family < 0xF) {
            host_model.model = base_model;
        } else {
            host_model.model = (ext_model << 4) | base_model;
        }

        host_model.stepping = eax & STEPPING_MASK;

        printf("family %u model %u stepping %d\n", host_model.family, host_model.model, host_model.stepping);

        return &host_model;
    }

    return NULL;
}

uint64_t build_vcpu_model(uint32_t family, uint32_t model, uint32_t stepping)
{
    uint64_t val;

    val = 0;
    // Stepping
    val |= stepping & 0xf;
    // BaseModel
    val |= (model & 0xf) << 4;
    // BaseFamily
    val |= (family > 0xf ? 0xf : family) << 8;
    // ExtModel
    val |= ((model >> 4) & 0xf) << 16;
    // ExtFamily
    val |= (family > 0xf ? (family - 0xf) : 0) << 20;

    return val;
}

uint64_t get_vcpu_model_value(const char* name)
{
    const struct x86_vcpu_model *model;
    uint64_t val;

    model = get_vcpu_model(name);
    if (!model)
        return 0x0ull;

    val = build_vcpu_model(model->family, model->model, model->stepping);

    return val;
}

/**
 * @brief set the vcpu register to reset state
 *
 * @param  [in]  buf                 The vmsa buffer
 * @param  [in]  buf_len             The buffer length
 * @param  [in]  model_name          The vcpu model
 * @param  [in]  reset_addr          The AP reset addr
 * @param  [in]  reset_data_valid    AP flag
 *
 * @returns 0 when success, -1 otherwise
 */
int vcpu_reset(void *buf, size_t buf_len,
               uint64_t vcpu_model_value,
               uint32_t reset_addr, bool reset_data_valid)
{
    csv_save_area *vmsa = (csv_save_area *)buf;

    if (!buf || buf_len != VMSA_SIZE ) {
        printf("Error: Invalid VMSA buffer :"
               " buf:0x%016lx buf_len:0x%x \n",
               (uint64_t)buf, (uint32_t)buf_len);
        return -1;
    }

    memset(buf, 0, buf_len);

    /* ES, CS, SS, DS, FS, GS */
    vmsa->es.selector = 0x0;
    // Qemu override @Accessed bit.
    vmsa->es.attrib = 0x92 | 0x1;
    vmsa->es.limit = 0xffffu;
    vmsa->es.base = 0x0ull;

    vmsa->cs.selector = 0xf000;
    // Qemu override @Accessed bit.
    vmsa->cs.attrib = 0x9a | 0x1;
    vmsa->cs.limit = 0xffffu;
    if (!reset_data_valid) {
        // BSP's CS.base follows the spec
        vmsa->cs.base = 0xffff0000ull;
    } else {
        // AP's CS.base is defined in guest's uefi bios
        vmsa->cs.base = (uint64_t)(reset_addr & 0xffff0000);
    }

    vmsa->ss.selector = 0x0;
    // Qemu override @Accessed bit.
    vmsa->ss.attrib = 0x92 | 0x1;
    vmsa->ss.limit = 0xffffu;
    vmsa->ss.base = 0x0ull;

    vmsa->ds.selector = 0x0;
    // Qemu override @Accessed bit.
    vmsa->ds.attrib = 0x92 | 0x1;
    vmsa->ds.limit = 0xffffu;
    vmsa->ds.base = 0x0ull;

    vmsa->fs.selector = 0x0;
    // Qemu override @Accessed bit.
    vmsa->fs.attrib = 0x92 | 0x1;
    vmsa->fs.limit = 0xffffu;
    vmsa->fs.base = 0x0ull;

    vmsa->gs.selector = 0x0;
    // Qemu override @Accessed bit.
    vmsa->gs.attrib = 0x92 | 0x1;
    vmsa->gs.limit = 0xffffu;
    vmsa->gs.base = 0x0ull;

    /* GDTR, LDTR, IDTR, TR */
    vmsa->gdtr.limit = 0xffffu;
    vmsa->gdtr.base = 0x0ull;

    vmsa->ldtr.selector = 0x0;
    vmsa->ldtr.attrib = 0x82;
    vmsa->ldtr.limit = 0xffffu;
    vmsa->ldtr.base = 0x0ull;

    vmsa->idtr.limit = 0xffffu;
    vmsa->idtr.base = 0x0ull;

    vmsa->tr.selector = 0x0;
    // Qemu override bit3 of Type field.
    vmsa->tr.attrib = 0x83 | 0x8;
    vmsa->tr.limit = 0xffffu;
    vmsa->tr.base = 0x0ull;

    /* PL0_SSP, PL1_SSP, PL2_SSP, PL3_SSP */
    /* U_CET, RESERVED, VMPL */

    /* CPL */
    vmsa->cpl = 0x0;

    /* RESERVED */

    /* EFER */
    // SVME bit must be set, VMRUN will check this bit.
    vmsa->efer = 0x0ull | (1 << 12);

    /* RESERVED */

    /* XSS */
    vmsa->xss = 0x0ull;

    /* CR4, CR3, CR0 */
    // KVM always set CR4.MCE bit.
    vmsa->cr4 = 0x0ull | (1 << 6);

    vmsa->cr3 = 0x0ull;

    // KVM always clear CR0.CD and CR0.NW bits to speed-up boot.
    vmsa->cr0 = 0x0000000060000010ull & ~(0x3ull << 29);

    /* DR7, DR6 */
    vmsa->dr7 = 0x0000000000000400ull;
    vmsa->dr6 = 0x00000000ffff0ff0ull;

    /* RFLAGS */
    vmsa->rflags = 0x0000000000000002ull;

    /* RIP */
    if (!reset_data_valid) {
        // BSP's CS.base follows the spec
        vmsa->rip = 0x000000000000fff0ull;
    } else {
        // AP's CS.base is defined in guest's uefi bios
        vmsa->rip = (uint64_t)(reset_addr & 0x0000ffff);
    }

    /* DR0, DR1, DR2, DR3 */
    /* DR0_ADDR_MASK, DR1_ADDR_MASK, DR2_ADDR_MASK, DR3_ADDR_MASK */

    /* RESERVED */

    /* RSP */
    vmsa->rsp = 0x0ull;

    /* S_CET, SSP, ISST_ADDR */

    /* RAX */
    vmsa->rax = 0x0ull;

    /* STAR, LSTAR, CSTAR, SFMASK, KERNEL_GS_BASE */
    vmsa->star = 0x0ull;
    vmsa->lstar = 0x0ull;
    vmsa->cstar = 0x0ull;
    vmsa->sfmask = 0x0ull;
    vmsa->kernel_gs_base = 0x0ull;

    /* SYSENTER_CS, SYSENTER_ESP, SYSENTER_EIP */
    vmsa->sysenter_cs = 0x0ull;
    vmsa->sysenter_esp = 0x0ull;
    vmsa->sysenter_eip = 0x0ull;

    /* CR2 */
    vmsa->cr2 = 0x0ull;

    /* RESERVED */

    /* G_PAT */
    vmsa->g_pat = 0x0007040600070406ull;

    /* DBGCTL, BR_FROM, BR_TO, LASTEXCPFROM, LASTEXCPTO */
    vmsa->dbgctl = 0x0ull;
    vmsa->br_from = 0x0ull;
    vmsa->br_to = 0x0ull;
    vmsa->last_excp_from = 0x0ull;
    vmsa->last_excp_to = 0x0ull;

    /* RESERVED */

    /* PKRU, TSC_AUX */

    /* RESERVED */

    /* RCX */
    vmsa->rcx = 0x0ull;

    /* RDX */
    // RDX should be the value of CPUID_Fn0000_0001_EAX
    vmsa->rdx = vcpu_model_value;

    /* RBX */
    vmsa->rbx = 0x0ull;

    /* RESERVED */

    /* RBP, RSI, RDI, R8, R9, R10, R11, R12, R13, R14, R15 */
    vmsa->rbp = 0x0ull;
    vmsa->rsi = 0x0ull;
    vmsa->rdi = 0x0ull;
    vmsa->r8 = 0x0ull;
    vmsa->r9 = 0x0ull;
    vmsa->r10 = 0x0ull;
    vmsa->r11 = 0x0ull;
    vmsa->r12 = 0x0ull;
    vmsa->r13 = 0x0ull;
    vmsa->r14 = 0x0ull;
    vmsa->r15 = 0x0ull;

    /* RESERVED */

    /* GUEST_EXITINFO1, GUEST_EXITINFO2, GUEST_EXITINTINFO, GUEST_NRIP */
    /* SEV_FEATURES, VINTR_CTRL */
    /* GUEST_EXITCODE, VIRTUAL_TOM */
    /* LTB_ID, PCPU_ID */
    /* EVENTINJ */

    /* XCR0 */
    vmsa->xcr0 = 0x0000000000000001ull;

    /* RESERVED */

    /* X87_DP, MXCSR, X87_FTW, X87_FSW */
    // Currently, KVM won't init these regs to reset state
    vmsa->x87_dp = 0x0ull;
    vmsa->mxcsr = 0x0; // 0x1f80
    vmsa->x87_ftw = 0x0; // 0x5555
    vmsa->x87_fsw = 0x0;

    /* X87_FCW, X87_FOP, X87_DS, X87_CS, X87_RIP */
    // Currently, KVM won't init these regs to reset state
    vmsa->x87_fcw = 0x0; // 0x0040
    vmsa->x87_fop = 0x0;
    vmsa->x87_ds = 0x0;
    vmsa->x87_cs = 0x0;
    vmsa->x87_rip = 0x0ull;

    /* FPREG_X87, FPREG_XMM, FPREG_YMM */
    memset(vmsa->fpreg_x87, 0, sizeof(vmsa->fpreg_x87));
    memset(vmsa->fpreg_xmm, 0, sizeof(vmsa->fpreg_xmm));
    memset(vmsa->fpreg_ymm, 0, sizeof(vmsa->fpreg_ymm));

    return 0;
}

/**
 * @brief  reverse data endian
 *
 * @param  [in]  bytes     The input data
 * @param  [in]  size      The input data length
 *
 */
void reverse_bytes(uint8_t* bytes, size_t size)
{
    uint8_t* start = bytes;
    uint8_t* end   = bytes + size - 1;

    while (start < end) {
        uint8_t byte = *start;
        *start       = *end;
        *end         = byte;
        start++;
        end--;
    }
}

/**
 * @brief  dump data in hex format
 *
 * @param  [in]  addr     The memory address
 * @param  [in]  len      The data length
 *
 */
void hex_dump(const uint8_t* addr, uint32_t len)
{
    int i;

    for (i = 0; i < len; i++) {
        if (i % 16 == 15) {
            printf(" %02X\n", *addr);
        } else {
            printf(" %02X", *addr);
        }
        addr++;
    }
    if (i % 16 != 0) {
        printf("\n");
    }
}

/**
 * @brief  convert string to integer
 *
 * @param  [in]  str     The input string
 *
 * @returns integer
 */
long str2int(char * str)
{
    long v = 0;

    if (!str)
        return 0;

    if (strlen(str) >= 2 && str[0] == '0' &&
       (str[1] == 'x' || str[1] == 'X')) {
        v = strtol(str, NULL, 16);
    } else {
        v = strtol(str, NULL, 10);
    }


    return v;
}



static const char base64_table[] =
    "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";

/**
 * encode array to base64
 * @param data  input to be encode
 * @param length input length
 * @return base64 data
 */
char* base64_encode(const unsigned char* data, size_t length) {
    size_t encoded_length = 0;
    char* encoded_data = NULL;
    size_t j = 0;
    size_t padding = 0;

    if (!data || length == 0)
        return NULL;

    encoded_length = 4 * ((length + 2) / 3);
    encoded_data = (char*)malloc(encoded_length + 1);
    if (!encoded_data)
        return NULL;


    for (size_t i = 0; i < length; i += 3) {
        uint32_t octet_a = i + 0 < length ? (unsigned char)data[i + 0] : 0;
        uint32_t octet_b = i + 1 < length ? (unsigned char)data[i + 1] : 0;
        uint32_t octet_c = i + 2 < length ? (unsigned char)data[i + 2] : 0;

        uint32_t triple = (octet_a << 0x10) + (octet_b << 0x08) + octet_c;

        encoded_data[j++] = base64_table[(triple >> 18) & 0x3F];
        encoded_data[j++] = base64_table[(triple >> 12) & 0x3F];
        encoded_data[j++] = base64_table[(triple >> 6)  & 0x3F];
        encoded_data[j++] = base64_table[triple        & 0x3F];
    }

    // Add padding
    padding = (3 - (length % 3)) % 3;
    for (size_t i = 0; i < padding; i++) {
        encoded_data[encoded_length - padding + i] = '=';
    }

    encoded_data[encoded_length] = '\0'; // Null-terminate the string

    return encoded_data;
}

